package com.ti.demo.mcp;

import io.modelcontextprotocol.server.McpSyncServerExchange;
import io.modelcontextprotocol.spec.McpSchema;
import org.springaicommunity.mcp.annotation.McpProgressToken;
import org.springaicommunity.mcp.annotation.McpTool;
import org.springaicommunity.mcp.annotation.McpToolParam;
import org.springframework.stereotype.Component;

import java.util.Map;

/**
 * @description: @McpTool
 * The @McpTool annotation marks a method as an MCP tool implementation with automatic JSON schema generation.
 * @Author Triagen
 * @Date 2025/9/29 16:49
 */
@Component
public class McpTools {

    // Basic Usage
    @McpTool(name = "add", description = "Add two numbers together")
    public int add(
            @McpToolParam(description = "First number", required = true) int a,
            @McpToolParam(description = "Second number", required = true) int b) {
        return a + b;
    }


    // Advanced Features
    public record AreaResult(double v, String square_units) {
    }

    @McpTool(name = "calculate-area",
            description = "Calculate the area of a rectangle",
            annotations = @McpTool.McpAnnotations(
                    title = "Rectangle Area Calculator",
                    readOnlyHint = true,
                    destructiveHint = false,
                    idempotentHint = true
            ))
    public AreaResult calculateRectangleArea(
            @McpToolParam(description = "Width", required = true) double width,
            @McpToolParam(description = "Height", required = true) double height) {

        return new AreaResult(width * height, "square units");
    }

    // With Server Exchange
    // Tools can access the server exchange for advanced operations:
    @McpTool(name = "process-data", description = "Process data with server context")
    public String processData(
            McpSyncServerExchange exchange,
            @McpToolParam(description = "Data to process", required = true) String data) {

        // Send logging notification
        exchange.loggingNotification(McpSchema.LoggingMessageNotification.builder()
                .level(McpSchema.LoggingLevel.INFO)
                .data("Processing data: " + data)
                .build());

        // Send progress notification if progress token is available
        exchange.progressNotification(new McpSchema.ProgressNotification(
                "progressToken", 0.5, 1.0, "Processing..."));

        return "Processed: " + data.toUpperCase();
    }

    // Dynamic Schema Support
    // Tools can accept CallToolRequest for runtime schema handling:
    @McpTool(name = "flexible-tool", description = "Process dynamic schema")
    public McpSchema.CallToolResult processDynamic(McpSchema.CallToolRequest request) {
        Map<String, Object> args = request.arguments();

        // Process based on runtime schema
        String result = "Processed " + args.size() + " arguments dynamically";

        return McpSchema.CallToolResult.builder()
                .addTextContent(result)
                .build();
    }

    // Progress Tracking
    // Tools can receive progress tokens for tracking long-running operations:
    @McpTool(name = "long-task", description = "Long-running task with progress")
    public String performLongTask(
            @McpProgressToken String progressToken,
            @McpToolParam(description = "Task name", required = true) String taskName,
            McpSyncServerExchange exchange) {

        if (progressToken != null) {
            exchange.progressNotification(new McpSchema.ProgressNotification(
                    progressToken, 0.0, 1.0, "Starting task"));

            // Perform work...

            exchange.progressNotification(new McpSchema.ProgressNotification(
                    progressToken, 1.0, 1.0, "Task completed"));
        }

        return "Task " + taskName + " completed";
    }
}
